import torch
import torch.nn as nn

from transcoder.models.block import TRANS_BLOCK

class DemodulateModel(nn.Module):
    def __init__(self, Nc, Q, patches):
        super().__init__()
        self.Q = Q
        self.Nc = Nc
        self.inEmbed = TRANS_BLOCK(2*Q, 2*Q, 1024, 6)
        self.adaptor = nn.Sequential(
            nn.Linear(2*self.Q, 512),
            nn.GELU(),
            nn.Linear(512, patches)
        )
        
    def forward(self, z): # [B, Nc, K, 2*Q]
        z = self.inEmbed(z)                  # [B*K, Nc, 2*Q]
        z = z.reshape(-1, self.Nc, 2*self.Q) # [B*K, Nc, 2*Q]
        z = self.adaptor(z)                  # [B*K, Nc, L]
        
        return z